import rsa
import time
import os.path

def LICENSE():
    print('''
        easyRSA - It just encrypt and decrypt.
        Copyright (C) 2021  stak

        This program is free software: you can redistribute it and/or modify
        it under the terms of the GNU General Public License as published by
        the Free Software Foundation, either version 3 of the License, or
        (at your option) any later version.

        This program is distributed in the hope that it will be useful,
        but WITHOUT ANY WARRANTY; without even the implied warranty of
        MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.   See the
        GNU General Public License for more details.

        You should have received a copy of the GNU General Public License
        along with this program.  If not, see <https://www.gnu.org/licenses/>.
    ''')

def console_picture():
    print('                      ____  ____    _     ')
    time.sleep(0.1)
    print('  ___  __ _ ___ _   _|  _ \/ ___|  / \    ')
    time.sleep(0.1)
    print(' / _ \/ _` / __| | | | |_) \___ \ / _ \   ')
    time.sleep(0.1)
    print('|  __/ (_| \__ \ |_| |  _ < ___) / ___ \  ')
    time.sleep(0.1)
    print(' \___|\__,_|___/\__, |_| \_\____/_/   \_\ ')
    time.sleep(0.1)
    print('                |___/                     ')

def gen_rsa(bit=2048):
    global pubkey, privkey
    if __name__ == '__main__':
        bit = input('Enter number of bits for generate RSA keys without spaces(default 2048): ')
        if bit == '':
            bit = 2048
        bit = int(bit)
    (pubkey, privkey) = rsa.newkeys(bit)
    pubkey_pem = pubkey.save_pkcs1()
    privkey_pem = privkey.save_pkcs1()
    if os.path.isfile('private.pem') == False:
        file_out = open("privkey.pem", "wb")
        file_out.write(privkey_pem)
        file_out.close()
    if os.path.isfile('public.pem') == False:
        file_out = open("public.pem", "wb")
        file_out.write(pubkey_pem)
        file_out.close()
    else:
        print('File(s) already exists!')

def encryption(word):
    global encrypted_message
    word = word.encode('utf-8')
    encrypted_message = rsa.encrypt(word, pubkey)
    return encrypted_message

def decryption(encrypted_message):
    unencrypted_message = rsa.decrypt(encrypted_message, privkey)
    global decode_message
    decode_message = unencrypted_message.decode('utf-8')
    return decode_message

def gen_rsa_q():
        userinput = input('Generate RSA keys?[y,n]: ')
        if userinput == 'y':
            gen_rsa()
        else:
            print('RSA keys have not been generated!')

def control():
    exit_words = ['exit', 'stop', 'break', '/exit', '/stop', '/break','-exit', '-stop', '-break`']
    while True:
        instruction = input('Enter instruction: ')
        if instruction == '-start':
            gen_rsa()
            message = input('Enter word: ')
            encryption(message)
            print(f'encrypted message = {encrypted_message}')
        if instruction == '-e':
            try:
                word = input('Enter word: ')
                encryption(word)
                print(f'encrypted message = {encrypted_message}')
            except NameError:
                print("You didn't generate RSA keys")
                gen_rsa_q()
            finally:
                continue

        if instruction == '-d':
            try:
                decryption(encrypted_message=encrypted_message)
                print(f'decode_message = {decode_message}')
            except NameError:
                print("You didn't encrypt any words, enter -e  to encrypt")
            finally:
                continue

        if instruction == '-g':
            gen_rsa()

        if instruction in exit_words:
            console_picture()
            print('Bye!')
            break

def main():
    LICENSE()
    console_picture()
    control()

if __name__ == "__main__":
    main()
